package main

import (
	"context"
	"fmt"
	chromago "github.com/amikos-tech/chroma-go"
	"github.com/amikos-tech/chroma-go/collection"
	collama "github.com/amikos-tech/chroma-go/ollama"
	"github.com/amikos-tech/chroma-go/types"
	"github.com/tmc/langchaingo/documentloaders"
	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/llms/ollama"
	"github.com/tmc/langchaingo/schema"
	"github.com/tmc/langchaingo/textsplitter"
	"hcy-api/lib/id"
	"os"
	"time"
)

func Rag(que string) (err error) {
	rs, err := SearchChunk(que)
	if err != nil {
		return err
	}
	qwen, err := ollama.New(
		ollama.WithServerURL("http://117.72.38.226:11434"),
		ollama.WithModel("qwen:4b"),
	)
	if err != nil {
		return
	}
	msgs := make([]llms.MessageContent, 0)
	systemDoc := make([]llms.ContentPart, 0)
	for i := range rs.Documents {
		systemDoc = append(systemDoc, llms.TextContent{Text: rs.Documents[i][0]})
	}
	msgs = append(msgs, llms.MessageContent{
		Role:  llms.ChatMessageTypeSystem,
		Parts: systemDoc,
	})
	userQ := make([]llms.ContentPart, 0)
	userQ = append(userQ, llms.TextContent{Text: que})
	msgs = append(msgs, llms.MessageContent{
		Role:  llms.ChatMessageTypeGeneric,
		Parts: userQ,
	})
	rsp, err := qwen.GenerateContent(
		context.Background(),
		msgs,
	)
	if err != nil {
		return err
	}
	for i := range rsp.Choices {
		fmt.Println(rsp.Choices[i].Content)
	}
	return
}

func SearchChunk(que string) (rs *chromago.QueryResults, err error) {
	client, err := chromago.NewClient("http://117.72.38.226:8000")
	if err != nil {
		return
	}
	ef, err := collama.NewOllamaEmbeddingFunction(
		collama.WithBaseURL("http://117.72.38.226:11434"),
		collama.WithModel("milkey/m3e"),
	)
	coll, err := client.GetCollection(
		context.Background(),
		"go-chroma",
		ef,
	)
	if err != nil {
		return
	}
	rs, err = coll.Query(context.Background(), []string{que}, 1, nil, nil, nil)
	return
}

func Embedding(fp string) (err error) {
	//m3e, err := ollama.New(
	//	ollama.WithServerURL("http://117.72.38.226:11434"),
	//	ollama.WithModel("milkey/m3e"),
	//)
	if err != nil {
		return err
	}
	docs, err := ReadPDF(fp)
	//texts := make([]string, len(docs))
	//for i := range docs {
	//	texts[i] = docs[i].PageContent
	//}
	////ctx := context.Background()
	////embs, err := m3e.CreateEmbedding(ctx, texts)
	////if err != nil {
	////	return err
	////}
	client, err := chromago.NewClient("http://117.72.38.226:8000")
	if err != nil {
		return err
	}

	ef, err := collama.NewOllamaEmbeddingFunction(
		collama.WithBaseURL("http://117.72.38.226:11434"),
		collama.WithModel("milkey/m3e"),
	)
	li, err := client.ListCollections(context.Background())
	if err != nil {
		return err
	}
	var coll *chromago.Collection
	for i := range li {
		if li[i].Name == "go-chroma" {
			coll = li[i]
		}
	}
	if coll == nil {
		coll, err = client.NewCollection(
			context.Background(),
			collection.WithEmbeddingFunction(ef),
			collection.WithName("go-chroma"),
			collection.WithHNSWDistanceFunction(types.L2), // 图索引，类似数据库的索引，后面还要具体了解这个算法
		)
		if err != nil {
			return err
		}
	}

	// 创建一个record 用于记录数据
	rs, err := types.NewRecordSet(
		types.WithEmbeddingFunction(ef),
		types.WithIDGenerator(types.NewULIDGenerator()),
	)
	if err != nil {
		return err
	}
	for i := range docs {
		rs.WithRecord(
			types.WithDocument(docs[i].PageContent),
			types.WithID(fmt.Sprint(id.GetSnowId())),
			types.WithMetadata("page", docs[i].Metadata["page"]),
			types.WithMetadata("total", docs[i].Metadata["total_pages"]),
		)
	}
	fmt.Println("前", rs)
	t1 := time.Now().Unix()
	// 这个方法命名有问题，这里实际执行了 embedding 和 校验两个步骤
	r, err := rs.BuildAndValidate(context.TODO())
	if err != nil {
		return err
	}
	t2 := time.Now().Unix()
	fmt.Println("后", r)
	fmt.Println(t2 - t1)
	// 写入数据
	_, err = coll.AddRecords(context.Background(), rs)
	if err != nil {
		return err
	}
	return
}

func ReadPDF(fp string) (docs []schema.Document, err error) {
	f, err := os.Open(fp)
	if err != nil {
		return
	}
	finfo, err := f.Stat()
	p := documentloaders.NewPDF(f, finfo.Size())
	split := textsplitter.NewRecursiveCharacter()
	split.ChunkSize = 1000   // size of the chunk is number of characters
	split.ChunkOverlap = 300 // overlap is the number of characters that the chunks overlap
	docs, err = p.LoadAndSplit(context.Background(), split)
	return
}
